import os
import math
import pickle
import numpy
from scipy import interpolate
from CalTorque import OpenCalTorqueXLS, GetSheetData
from scipy.optimize import minimize
from scipy.optimize import curve_fit
import matplotlib
from matplotlib import pyplot as PLT

#____________________________________________________________________________

def GetRunData(path = 'data/'):
    datapath = path
    lsdir = os.listdir(datapath)

    for f in lsdir:
        if (f.split('.')[-1] not in ['xlsx', 'xls']):
            continue
    print f
    book = OpenCalTorqueXLS(datapath+f)
    data = GetSheetData(book)
    
    return data, f


def GetTestRunData(path):
    book = OpenCalTorqueXLS(path)
    data = GetSheetData(book)
    return data
    
    
def GetPklAverage(path = 'pickled/pickled_average_in.pkl'):
    file = open(path, 'rb')
    avgdata = pickle.load(file)
    file.close()

    #Xavg = numpy.array(avgdata['x'])    #len(Xavg) = 400
    #Yavg = numpy.array(avgdata['y'])
    
    return avgdata

def AverageTorqueInterp(x, offset = 0):            ### if put x as argument, would there be an offset arg too? OFFSET FOR AVERAGE
    x2 = []
    f = interpolate.interp1d(Xavg, Yavg)
    if offset != 0:
        for i, j in enumerate(x):
            x2.append(x[i] - offset)
    else:
        x2 = x

    #Xin = []
    #for i in x2:
    #    if i > 10 and i < 390:
    #        Xin.append(i)
    #print len(Xin)
    return f(x2)


def PolyTwoGauss(p, x):
    offset = 52            ### distance between peaks
    #p[0],p[1],p[2] = amplitude , center_x, sigma: for peak A
    #p[3],p[4] = amplitude, sigma: for peak B
    #p[5],p[6],p[7] = ax^2 + bx + c
    x0 = x[0]
    y = numpy.zeros(len(x0))
    for i, j in enumerate(x0):  #index 1 is out of bounds for axis 0 with size 1
        y[i] = p[0]*math.exp(-0.5*((j-p[1])/p[2])**2.) + p[3]*math.exp(-0.5*((j-(p[1]+offset))/p[4])**2.) + p[6]*(j - p[5])**2 + p[7]*(j - p[5]) + p[8]
        if y[i] <= 1e-5: 
            y[i] = 1e-5
    return y


def PolyFn(x, x0, a, b, c):     # assumes ydata = f(xdata, params)
   return a*(x-x0)**2 + b*(x-x0) + c

def Residual(p, y, x):
    modely = PolyTwoGauss(p, x)
    return numpy.sum((modely - y)**2)


### idk if binning the data is the best option - maybe just add something after it?
def BinData(x, y):
    bins = 200.0
    max_d = 400.0
    x_inputs = numpy.zeros(bins)
    xdata = numpy.linspace(0, max_d, bins)
    ydata = numpy.zeros(bins)

    for j, k in enumerate(xdata):
        pos_torque = []
        x0 = k - (max_d / bins) / 2.
        x1 = k + (max_d / bins) / 2.
        for i, p in enumerate(x):
           if p > x0 and p <= x1:
              pos_torque.append(y[i])
        #print len(pos_torque), k
        if len(pos_torque) > 0:
            ydata[j] += max(pos_torque)
            x_inputs[j] += 1.

    for i, n in enumerate(x_inputs):
        if n == 0:
            continue
        ydata[i] = ydata[i]/n
    return xdata, ydata


#delta - minimum difference between peak and following pts
def PeakFinder1(y_ax, x_ax, lookahead = 3, delta = 1):
    #delta - minimum difference between peak and following pts
    max_peaks = {}
    max_peaks['maxpos'] = []
    max_peaks['max'] = []
    peaklocation = []
    max = y_ax[(numpy.abs(x_ax - 200)).argmin()]    # this starts search for peaks at 100th bin (x_ax = 200) - also script won't fail if there is a really big initial torque

    # Get peaks
    for i, y in enumerate(y_ax[: -lookahead]):
        if y > max:
            max = y     #maxpos referenced before assignment.....y <= max, thus maxpos does not exist
            maxpos = x_ax[i]  ### There has to be a better way of doing this.
            
            if y < max - delta and max != numpy.Inf:    # I think it might be a bad idea to layer this many ifs
                if y_ax[i:i + lookahead].max() < max:
                    if maxpos not in max_peaks['maxpos']:   # for some reason the previous code returns multiple copies of the same peak - this prevents that
                        max_peaks['maxpos'].append(maxpos)
                        max_peaks['max'].append(max)
    print max_peaks['maxpos']
    # Filter peaks to only the ones in which we are interested
    for i, j in enumerate(max_peaks['maxpos']):
        if 300 <= j < 330:
            peaklocation.append(j)
        elif 350 < j <= 380:
            peaklocation.append(j)
    # maybe add something that raises error if there are more than 2 peaks?
    
    return peaklocation


def PeakFinder2(y_ax, x_ax):
    torquebtw300_330 = []
    torquebtw350_380 = []
    min_before = [0, 0]
    min_before_pos = [0, 0]
    min_after = [0, 0]
    min_after_pos = [0, 0]
    
    for i, k in enumerate(x_ax):
        if 300 <= k < 330:
            torquebtw300_330.append(y_ax[i])
            max1 = max(torquebtw300_330)
            maxpos1 = x_ax[list(y_ax).index(max1)]
        elif 350 < k <= 375:
            torquebtw350_380.append(y_ax[i])
            max2 = max(torquebtw350_380)
            maxpos2 = x_ax[list(y_ax).index(max2)]
    
    for m, n in enumerate(x_ax):
        if n == maxpos1:
            min_before[0] = min(y_ax[m - 5 : m])
            min_before_pos[0] = x_ax[list(y_ax).index(min_before[0])]
            min_after[0] = min(y_ax[m : m + 10])
            min_after_pos[0] = x_ax[list(y_ax).index(min_after[0])]
        elif n == maxpos2:
            min_before[1] = min(y_ax[m - 5 : m])
            min_before_pos[1] = x_ax[list(y_ax).index(min_before[1])]
            min_after[1] = min(y_ax[m : m + 10])
            min_after_pos[1] = x_ax[list(y_ax).index(min_after[1])]
    
    B2peak = min_before_pos[0] + ((min_after_pos[0] - min_before_pos[0])/2)
    B3peak = min_before_pos[1] + ((min_after_pos[1] - min_before_pos[1])/2)

    return B2peak, B3peak


def CorrectForOffset_P(xdata, ydata):
    
    Xavg = numpy.array(GetPklAverage()['x'])
    Yavg = numpy.array(GetPklAverage()['y'])

    x = Xavg[140:-10]
    y = Yavg[140:-10]
    
    poly_fit = curve_fit(PolyFn, x, y) #to get initial guess on the baseline
    poly_parms = poly_fit[0]

    # put in initial guesses now
    fit_parms = [5, 320, 5, 5, 5, poly_parms[0], poly_parms[1], poly_parms[2], poly_parms[3]]
    fit_bounds = [(0,20),(300,350), (0, 40), (0, 20), (0, 20), (None, None), (None,None), (None, None), (None, None)]

    result = minimize(Residual, fit_parms, args = (([y]), ([x])), method='TNC', bounds=fit_bounds,tol=0.0000005,options={'disp':False})
    new_parms = result.x
    
    # gets peaks from binned data
    bindata = BinData(xdata, ydata)
    peaklocation = PeakFinder2(bindata[1], bindata[0])
    
    #print peaklocation
    offset1 = new_parms[1] - peaklocation[0]
    offset2 = (new_parms[1] + 52) - peaklocation[1]
    offset = (offset1 + offset2)/2
    
    ### offsets are different, which one use? or use an average?
    
    # produces position data array without the offset
    offset_less_pos = []
    for i, k in enumerate(xdata):
        thing = k + offset    
        offset_less_pos.append(thing)
    
    return offset_less_pos, offset, new_parms


#_______________________________________________________

def PlotFit_peak(data, outdir = 'fitplots/peak/'):

    offset_less, offset, new_parms = CorrectForOffset_P(data['position'], data['torque_inoz'])

    run  = data['run']
    direction = data['direction']
    
    bindata = BinData(data['position'], data['torque_inoz'])
    x = bindata[0][140:-10]
    y = bindata[1][140:-10]
    
    fig = PLT.figure(figsize=(15,8),dpi=150)
    box = [0.14, 0.14, 0.76, 0.76]
    ax1 = fig.add_axes(box)
    ax1.set_ylabel('Torque (in.oz)')
    ax1.set_xlabel('Position (cm)')
    maxy = max(data['torque_inoz'])
    maxx = max(data['position'])
    miny = min(data['torque_inoz'])
    
    ## Add station locations
    if maxx >= 328.9:
        ax1.axvline(338.9, color = 'r', linewidth = 1, linestyle = '-')
        ax1.annotate('S2', xy = (340.9, miny * 0.8), color = 'r')
    if maxx >= 382.6:
        ax1.axvline(392.6, color = 'r', linewidth = 1, linestyle = '-')
        ax1.annotate('S5', xy = (394.6, miny * 0.8), color = 'r')
    if maxx >= 435.7:
        ax1.axvline(445.7, color = 'r', linewidth = 1, linestyle = '-')
        ax1.annotate('S8', xy = (447.7, miny * 0.8), color = 'r')
    if maxx >= 491.7:
        ax1.axvline(501.7, color = 'r', linewidth = 1, linestyle = '-')
        ax1.annotate('S11', xy = (503.7, miny * 0.8), color = 'r')
    if maxx >= 595.8:
        ax1.axvline(605.8, color = 'r', linewidth = 1, linestyle = '-')
        ax1.annotate('S17', xy = (607.8, miny * 0.8), color = 'r')
    
    ## Add bend locations
    if maxx >= 47.8:
        ax1.axvspan(47.8, 93.5, color = 'g', alpha = 0.1)
        ax1.annotate('Bellows', xy = (48, miny), color = 'g')
    if maxx >= 136.4:
        ax1.axvspan(136.4, 144.6, color = 'g', alpha = 0.1)
        ax1.annotate('First outer bend', xy = (136, miny), color = 'g')
    if maxx >= 196.5:
        ax1.axvspan(196.5, 218.0, color = 'g', alpha = 0.1)
        ax1.annotate('The "S"', xy = (212, miny), color = 'g')
    if maxx >= 228.3:
        ax1.axvspan(228.3, 241.5, color = 'g', alpha = 0.1)
    if maxx >= 274.6:
        ax1.axvspan(274.6, 284.6, color = 'g', alpha = 0.1)
        ax1.annotate('Bend 1', xy = (274, miny), color = 'g')
    if maxx >= 309.0:
        ax1.axvspan(309.0, 318.9, color = 'g', alpha = 0.1)
        ax1.annotate('Bend 2', xy = (309, miny), color = 'g')
    if maxx >= 358.8:
        ax1.axvspan(358.8, 370.8, color = 'g', alpha = 0.1)
        ax1.annotate('Bend 3', xy = (359, miny), color = 'g')
    if maxx >= 414.5:
        ax1.axvspan(414.5, 426.5, color = 'g', alpha = 0.1)
        ax1.annotate('Bend 4', xy = (414, miny), color = 'g')
    if maxx >= 445.3:
        ax1.axvspan(445.3, 455.3, color = 'g', alpha = 0.1)
        ax1.annotate('Bend 5', xy = (445, miny), color = 'g')
    if maxx >= 467.5:
        ax1.axvspan(467.5, 479.5, color = 'g', alpha = 0.1)
        ax1.annotate('Bend 6', xy = (467, miny), color = 'g')
    if maxx >= 523.8:
        ax1.axvspan(523.8, 535.8, color = 'g', alpha = 0.1)
        ax1.annotate('Bend 7', xy = (523, miny), color = 'g')
    if maxx >= 571.6:
        ax1.axvspan(571.6, 583.6, color = 'g', alpha = 0.1)
        ax1.annotate('Bend 8', xy = (571, miny), color = 'g')
    if maxx >= 627.9:
        ax1.axvspan(627.9, 639.9, color = 'g', alpha = 0.1)
        ax1.annotate('Bend 9', xy = (627, miny), color = 'g')
    if maxx >= 652.7:
        ax1.axvspan(652.7, 660.7, color = 'g', alpha = 0.1)
        ax1.annotate('Bend 10', xy = (652, miny), color = 'g')
    if maxx >= 678.5:
        ax1.axvspan(678.5, 681.1, color = 'g', alpha = 0.1)
        ax1.annotate('Bend 11', xy = (678, miny), color = 'g')

    PLT.title(run + ' Source ' + direction + ' Offset Fix by Peak Finder Method')
    ax1.grid()
    ax1.plot(x, y, '-', color = 'maroon', linewidth = 1, label = 'Binned Original Data')
    #ax1.plot(x, PolyTwoGauss(new_parms, [x]), '-', color = 'g', linewidth = 1, label = 'Poly-Two-Gauss')
    ax1.plot(offset_less, data['torque_inoz'], '-', color = 'blue', linewidth = 1, label = 'Corrected Offset')
    h, l = ax1.get_legend_handles_labels()
    PLT.legend(h, l, 'upper left')
    fig.savefig(outdir + run + '_' + direction + '_peakfindfit' + '.png')